import torch
import logging
import torch.nn as nn
from models.backbones.hrnet.seg_hrnet import get_hrnet



class RGBTCC(nn.Module):
    def __init__(self,):
        super(RGBTCC, self).__init__()

        self.last_layer = nn.Sequential(
            nn.Conv2d(
                in_channels=720,
                out_channels=720,
                kernel_size=1,
                stride=1,
                padding=0),
            nn.BatchNorm2d(720, momentum=0.01),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(720, 64, 4, stride=2, padding=1, output_padding=0, bias=True),
            nn.ReLU(inplace=True),
            nn.ConvTranspose2d(64, 1, 4, stride=2, padding=1, output_padding=0, bias=True),

        )
        
        self.init_weights()
        
        self.backbone = get_hrnet()
    
    def init_weights(self,):
        logging.info('=> init weights from normal distribution')
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.normal_(m.weight, std=0.001)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

    def forward(self, x):

        f = self.backbone(x)
        x = self.last_layer(f)

        return x


def get_model():
    rgbtcc = RGBTCC()
    return rgbtcc